#!/usr/bin/python3
# 从miranda的结果中提取需要的结果，计算ceRNA对
import sys
import os
import re
import fileinput
import multiprocessing
import ParametersError
import datetime
import cernaclass as cc
from mne.stats import fdr_correction
import scipy.stats as stats
from File import FileWriter


def about():
    print('\t\t\t\t\t本脚本')
    print('\t用来从miranda的输出结果中提取需要的配对结果')
    print('\t输出文件为.miranda结尾的utf-8编码文本文件\n')
    print('\t##########\t文档说明\t##########')
    print('\t-i：要输入的的文件，多个文件以逗号分隔')
    print('\t-o：要输出的文件')
    print('\t\t若-o参数为空，默认使用./cerna.txt')
    print('\t-p：要使用的进程数')
    exit()


class multi(multiprocessing.Process):
    def __init__(self, all,cerna, N, thread, output_file):
        multiprocessing.Process.__init__(self)
        self.all = all
        self.cerna = cerna
        self.total_num = N
        self.thread = thread
        self.output_file = output_file

    def run(self):
        all = self.all
        cerna = self.cerna
        total_num = self.total_num
        thread = self.thread
        output_file = self.output_file
        output_file = output_file + ".tem" + str(thread)
        count = 0
        jindu = 0
        total = len(cerna) * len(all)

        with FileWriter(output_file) as w:
            for i in range(len(all)):
                for j in range(len(cerna)):
                    count += 1
                    if all[i] & cerna[j]:
                        if all[i] == cerna[j]:
                            continue
                        bign = all[i] & cerna[j]
                        smalln = len(all[i].get_mirna())
                        smallk = len(cerna[j].get_mirna())
                        p = get_p(N=total_num, n=smalln, k=smallk, c=bign)

                        w.write(all[i].get_rna() + "\t")
                        # w.write(",".join(cerna[i].get_mirna()) + "\t")
                        w.write(cerna[j].get_rna() + "\t")
                        # w.write(",".join(cerna[j].get_mirna()) + "\t")
                        w.write(str(total_num) + "\t")
                        w.write(str(smallk) + "\t")
                        w.write(str(smalln) + "\t")
                        w.write(str(bign) + "\t")
                        w.write(str(p) + "\n")
                    if jindu != int((count/total)*100):
                        jindu = int((count / total) * 100)
                        print('进程' + str(thread+1) + '：计算进度' + str(jindu) + "%")


def dividFlags(flag, thread):                         # 用来分割的函数flags的函数
    if thread == 1:
        return list(flag)
    elif thread > len(flag):
        return dividFlags(flag, len(flag))             # 递归一下，防止进程数大于能需要被分配的量，无法正常分配
    elif thread > 1:
        print("分配任务中")
        results = []
        size = len(flag)//thread

        for i in range(thread):
            if i == (thread - 1):
                results.append(flag[i * size:])
            else:
                results.append(flag[i * size: size * (i+1)])
        return results


def factorial(N, n=1):
    if N > n:
        return N * factorial(N-1, n)
    else:
        return n


def cummin(N, n=1):
    if N > n:
        return N + cummin(N-1, n)
    else:
        return 1


def get_p(N, k, n, c):
    p = stats.hypergeom.sf(c, N+n, n, k)                         # scipy 中的phyper（R），为什么这么用，不知道，因为没看过官方文档，但是结果正确
    return p


def cerna_pairs(input_file, output_file, thread):
    begin = datetime.datetime.now()
    print("Find_ceRNA程序开始于：%s" % begin)

    pairs = {}
    mirna = set()
    if not isinstance(thread, int):
        thread = int(thread)

    if input_file:  # 判断一下输入的文件数量，分割一下，把文件处理好
        input_file = re.split(r',', input_file)
        for file in iter(input_file):
            if not os.path.isfile(file):
                print(file + "不是文件")
                exit()
    else:
        print("没有任何输入文件")
        exit()

    if not output_file:
        pattern = re.compile('\\\|/')
        end = pattern.search(input_file[1]).end()
        output_file = input_file[1][:end]
    else:
        ParametersError.raise_file(output_file)

    for i in input_file:
        with fileinput.input(i) as f:
            for line in f:
                lines = re.split("\\s+", line)
                mirna.add(lines[0])
                if lines[1] not in pairs.keys():
                    tem = set()
                    tem.add(lines[0])
                    pairs.update({lines[1]: tem})
                else:
                    tem = pairs.get(lines[1])
                    tem.add(lines[0])
                    pairs.update({lines[1]: tem})
    cerna = []
    for key, value in pairs.items():
        cerna.append(cc.CeRNAPairs(rna=key, mirna=value))

    pairs = None                                # 腾出内存空间

    if thread > 1:
        cerna_divided = dividFlags(cerna, thread)
        proc = []
        for i in range(thread):
            proc.append(multi(cerna=cerna_divided[i], N=len(mirna), thread=i, output_file=output_file, all=cerna))

        for i in range(thread):
            proc[i].start()

        for i in range(thread):
            proc[i].join()
    else:
        proc = multi(cerna=cerna, N=len(mirna), thread=0, output_file=output_file, all=cerna)
        proc.start()
        proc.join()
    cerna = []
    p = []
    for i in range(thread):
        count = 0
        input_tem = output_file + ".tem" + str(i)
        with fileinput.input(input_tem) as f:
            for line in f:
                count += 1
                line = line.strip("\n")                                 # 特别主要要去掉\n
                lines = re.split("\\s+", line)
                try:
                    p.append(float(lines[6]))
                except ValueError as e:
                    print(e.__traceback__)
                    print(input_tem)
                    print(lines[len(lines)-1])
                    print(count)
                    exit()
                cerna.append(line)
        os.remove(input_tem)
    q = fdr_correction(p)               # 用来算q值的
    value = list(q[1].tolist())
    diff = list(q[0].tolist())
    with FileWriter(output_file) as w:
        for i in range(len(cerna)):
            w.write(cerna[i] + "\t" + str(value[i]) + "\t" + str(diff[i]) + '\n')

    end = datetime.datetime.now()
    print("程序终止于：%s" % end)
    print("一共花费：%s" % str(end - begin))

if __name__ == '__main__':
    parameters = {}  # 用来读取命令行参数的，遍历所有参数，开头以-开始的，下一个就作为参数附加，然后统一添加到一个词典中，再提取出来
    if len(sys.argv) <= 1:
        about()
    try:
        for i in range(len(sys.argv)):
            if sys.argv[i] == '-h' or sys.argv[i] == '--help':
                about()
            if sys.argv[i].startswith('-'):
                parameters.update({sys.argv[i]: sys.argv[i + 1]})
            elif sys.argv[i].startswith('--'):
                parameters.update({sys.argv[i]: sys.argv[i + 1]})
    except IndexError:  # 判断参数是否输入有误
        print("输入的参数有误")
        exit()

    input_file = parameters.get('-i')  # 提取所有的输入文件

    output_file = parameters.get('-o')  # 提取所有的输入文件

    thread = parameters.get('-p')  # 提取进程数
    ParametersError.raise_int(thread)

    cerna_pairs(input_file, output_file, thread)
    input('任意键退出')
